package org.forum.jstun;

import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;

import org.forum.jstun.StunMessage.StunAddress;


public class NatDetector implements Runnable {
	public static final int NAT_ERROR				= 0;
	public static final int NAT_FULL_CONE 			= 1;
	public static final int NAT_RESTRICTED_CONE 	= 2;
	public static final int NAT_PORT_RESTRICTED_CONE= 3;
	public static final int NAT_SYMMETRIC			= 4;
	public static final int NAT_NONE_OPEN_INTERNET	= 10;
	public static final int NAT_NONE_WITH_FIREWALL	= 11;

	public static final int STATE_START				= 10;
	public static final int STATE_DETECT_SERVER		= 20;
	public static final int STATE_DETECT_FULL_CONE	= 30;
	public static final int STATE_DETECT_SYMMETRIC	= 40;
	public static final int STATE_DETECT_RESTRICTED	= 50;
	public static final int STATE_SUCCESS			= 100;
	public static final int STATE_FAILURE			= 101;
	public static final int ERROR_UDP_FAILURE		= 104;
	public static final int ERROR_SERVER_UNAVAILABLE= 404;
	public static final int ERROR_SERVER_ERROR		= 500;

	public interface DetectListener {
		public void onFinish(int state);
	}

	private String mLocalIp;
	private String mServerHost;
	private InetAddress mServerAddress;
	private int mServerPort;
	private InetAddress mChangedAddress;
	private int mChangedPort;
	private InetAddress mMappedAddress;
	private int mMappedPort;
	private String mUsername;
	private String mPassword;
	private int mState = ERROR_SERVER_UNAVAILABLE;
	private int mNatType = NAT_ERROR;
	private StunAttribute mErrorCode = null;

	private int mReceiveTimeout = 300;
	private int mDetectTimeout 	= 1000*15;
	private DatagramSocket mSocket;
	private DetectListener mListener = null;

	public NatDetector(String server, int port) {
		mServerHost = server;
		mServerPort = port;
	}

	public NatDetector(String server, int port, String localIp) {
		mServerHost = server;
		mServerPort = port;
		mLocalIp = localIp;
	}

	public void setTimeout(int receiveMillis, int detectMillis) {
		mReceiveTimeout = receiveMillis;
		mDetectTimeout = detectMillis;
	}

	public void setShareSecret(String username, String password) {
		mUsername = username;
		mPassword = password;
	}

	public void setDetectListener(DetectListener listener) {
		mListener = listener;
	}

	@Override
	public void run() {
		detect();
		if (mListener!=null) mListener.onFinish(mState);
	}

	public void stop() {
		if (mSocket!=null) mSocket.close();
		mSocket = null;
	}

	public int getState() {
		return mState;
	}

	public int getNatType() {
		return mNatType;
	}

	public StunAttribute getErrorCode() {
		return mErrorCode;
	}

	public InetAddress getMappedAddress() {
		return mMappedAddress;
	}

	public int getMappedPort() {
		return mMappedPort;
	}

	private void detect() {
		final int timespan = mDetectTimeout/3;
		mState = STATE_FAILURE;
		mNatType = NAT_ERROR;
		mSocket = openUdpSocket(mReceiveTimeout);
		if (mSocket==null) {
			mState = ERROR_UDP_FAILURE;
			return;
		}
		if (!hasStunServer(mSocket, timespan)) {
			mState = ERROR_SERVER_UNAVAILABLE;
			return;
		}
		if (mLocalIp!=null && !isUnderNat(mSocket, timespan)) {
			System.out.println("check open internet or firewall");
			if (isFullCone(mSocket, timespan)) mNatType = NAT_NONE_OPEN_INTERNET;
			else mNatType = NAT_NONE_WITH_FIREWALL;
			mState = STATE_SUCCESS;
			return;
		}
		System.out.println("check if is full cone nat");
		if (isFullCone(mSocket, timespan)) {
			mState = STATE_SUCCESS;
			mNatType = NAT_FULL_CONE;
			return;
		}
		System.out.println("check if is symmetric nat");
		if (isSymmetric(mSocket, timespan)) {
			mState = STATE_SUCCESS;
			mNatType = NAT_SYMMETRIC;
			return;
		}
		System.out.println("check restricted nat");
		if (isPortRestricted(mSocket, timespan)) {
			mState = STATE_SUCCESS;
			mNatType = NAT_PORT_RESTRICTED_CONE;
		} else {
			mState = STATE_SUCCESS;
			mNatType = NAT_RESTRICTED_CONE;
		}
	}

	private DatagramSocket openUdpSocket(int receiveTimeout) {
		DatagramSocket socket = null;
		try {
			mServerAddress = InetAddress.getByName(mServerHost);
			socket = new DatagramSocket();
			socket.setReuseAddress(true);
			socket.setSoTimeout(receiveTimeout);
			System.out.println("connect to server "+mServerAddress.getHostAddress()+":"+mServerPort);
		} catch (Exception e) {
			e.printStackTrace();
			if (socket!=null) socket.close();
			socket = null;
		}
		return socket;
	}

	private StunMessage buildStunMessage(int state) {
		StunMessage msg = new StunMessage(StunMessage.TYPE_BINDING_REQUEST);
		msg.setTransactionId(StunMessage.generateId(false));
		if (mUsername!=null && !mUsername.isEmpty()) {
			msg.appendAttribute(new StunMessage.UsernamePassword(
				StunAttribute.TYPE_USERNAME, mUsername));
		}
		switch (state) {
			case STATE_DETECT_SERVER:
			case STATE_DETECT_SYMMETRIC:
				msg.appendAttribute(new StunMessage.ChangeRequest(false, false));
				break;

			case STATE_DETECT_RESTRICTED:
				msg.appendAttribute(new StunMessage.ChangeRequest(false, true));
				break;

			case STATE_DETECT_FULL_CONE:
				msg.appendAttribute(new StunMessage.ChangeRequest(true, true));
				break;
		
			default:
				break;
		}
		if (mPassword!=null && !mPassword.isEmpty()) {
			msg.appendAttribute(StunMessage.generateMessageIntegrity(
				msg.getBytes(), mPassword.getBytes()));
		}
		return msg;
	}

	private StunMessage sendAndReceive(DatagramSocket socket, StunMessage msg, 
			int timespan, InetAddress destAddress, int destPort, boolean check) {
		byte[] data = msg.getBytes();
		DatagramPacket packet = new DatagramPacket(data, data.length);
		packet.setAddress(destAddress);
		packet.setPort(destPort);
		// stun over UDP without stream index, internet MTU recommend it is 576,
		// udp data length max is 548 (IP 20, UDP header 8).
		DatagramPacket receive = new DatagramPacket(new byte[1024], 1024);
		StunMessage response = null;
		for(int ts=0; response==null && ts<timespan; ts+=mReceiveTimeout) {
			try {
				socket.send(packet);
				while(true) {
					// already set socket SoTimeout
					socket.receive(receive);
					if (check && (!receive.getAddress().equals(packet.getAddress()) ||
								receive.getPort()!=packet.getPort()))
						continue;
					// parse receive message
					response = StunMessage.parseHeader(receive.getData());
					if (response==null) continue;
					if (StunMessage.transactionIdEquals(msg.getTransactionId(), 
							response.getTransactionId())) {
						response = StunMessage.parse(receive.getData(), receive.getLength());
						break;
					}
					response = null;
				}
			} catch (Exception e) {}
		}
		return response;
	}

	private boolean hasStunServer(DatagramSocket socket, int timespan) {
		// get mapped Address via stun server
		StunMessage msg = buildStunMessage(STATE_DETECT_SERVER);
		StunMessage response = sendAndReceive(socket, msg, timespan, 
				mServerAddress, mServerPort, true);
		if (response==null) return false;	// stun server unavailable
		// check stun message
		StunAttribute ec = response.getAttribute(StunAttribute.TYPE_ERROR_CODE);
		StunAttribute ma = response.getAttribute(StunAttribute.TYPE_MAPPED_ADDRESS);
		StunAttribute ca = response.getAttribute(StunAttribute.TYPE_CHANGED_ADDRESS);
		if (ma==null || ca==null || ec!=null) {
			// server error
			mErrorCode = ec;
			return false;
		}
		mMappedAddress = ((StunAddress) ma).getAddress();
		mMappedPort = ((StunAddress) ma).getPort();
		mChangedAddress = ((StunAddress) ca).getAddress();
		mChangedPort = ((StunAddress) ca).getPort();
		return true;
	}

	private boolean isUnderNat(DatagramSocket socket, int timespan) {
		// check mapped address equals to local address
		System.out.println("local "+mLocalIp+":"+socket.getLocalPort()
				+", mapped "+mMappedAddress+":"+mMappedPort);
		return !mMappedAddress.getHostAddress().equals(mLocalIp)
				|| socket.getLocalPort() != mMappedPort;
	}

	private boolean isFullCone(DatagramSocket socket, int timespan) {
		// send stun message with changeIp, changePort
		StunMessage msg = buildStunMessage(STATE_DETECT_FULL_CONE);
		StunMessage response = sendAndReceive(socket, msg, timespan,
				mServerAddress, mServerPort, false);
		return response!=null;	// full cone or not
	}

	private boolean isSymmetric(DatagramSocket socket, int timespan) {
		// get mapped address via changed address
		StunMessage msg = buildStunMessage(STATE_DETECT_SYMMETRIC);
		StunMessage response = sendAndReceive(socket, msg, timespan, 
				mChangedAddress, mChangedPort, true);
		if (response==null) return false;	// stun server unavailable
		// check stun message
		StunAttribute ec = response.getAttribute(StunAttribute.TYPE_ERROR_CODE);
		StunAddress ma = (StunAddress)response.getAttribute(StunAttribute.TYPE_MAPPED_ADDRESS);
		if (ma==null || ec!=null) {
			// server error
			mErrorCode = ec;
			return false;
		}
		return !mMappedAddress.equals(ma.getAddress()) || mMappedPort != ma.getPort();
	}

	private boolean isPortRestricted(DatagramSocket socket, int timespan) {
		// send stun message with changePort
		StunMessage msg = buildStunMessage(STATE_DETECT_RESTRICTED);
		StunMessage response = sendAndReceive(socket, msg, timespan,
				mServerAddress, mServerPort, false);
		return response==null;	// port restricted or not(ip restricted)
	}
}